# rabbitmq.py
import aio_pika
from typing import AsyncGenerator

# RabbitMQ 连接 URL
RABBITMQ_URL = "amqp://free2walk:free2walk@192.168.25.202:5672/"

_connection: aio_pika.abc.AbstractConnection | None = None
_channel: aio_pika.abc.AbstractChannel | None = None

async def get_rabbit_connection() -> aio_pika.abc.AbstractConnection:
    global _connection
    if _connection is None or not _connection.connected:
        _connection = await aio_pika.connect_robust(RABBITMQ_URL,
    heartbeat=60,
    timeout=10000)
    return _connection

async def get_rabbit_channel() -> aio_pika.abc.AbstractChannel:
    global _channel
    if _channel is None or not _channel.is_open:
        connection = await get_rabbit_connection()
        _channel = await connection.channel()
        # 可选：设置 QoS（公平分发）
        await _channel.set_qos(prefetch_count=1)
    return _channel

async def close_rabbitmq():
    global _channel, _connection
    if _channel:
        await _channel.close()
    if _connection:
        await _connection.close()